import os
from abc import ABC, abstractmethod
from collections import OrderedDict

import loguru
import torch

from utils.CommonUtils import get_current_timestamp, check_create_dir


class BaseModel(ABC):
    def __init__(self, config):
        self.logger = loguru.logger
        self.config = config
        self.model_save_dir = os.path.join(config.basic_dir, config.model_save_dir,
                                           get_current_timestamp() + "_" + config.name)
        if self.config.run_type == "train":
            check_create_dir(self.model_save_dir)
        self.model_names = list()
        self.loss_names = list()

    @abstractmethod
    def set_input(self, input):
        """Unpack input data from the dataloader and perform necessary pre-processing steps.

        Parameters:
            input (dict): includes the data itself and its metadata information.
        """
        pass

    @abstractmethod
    def forward(self):
        """Run forward pass; called by both functions <optimize_parameters> and <test>."""

        pass

    @abstractmethod
    def optimize_parameters(self):
        """Calculate losses, gradients, and update network weights; called in every training iteration"""
        pass

    @abstractmethod
    def backward(self):
        """Run forward pass; called by both functions <optimize_parameters> and <test>."""

        pass

    @abstractmethod
    def update_learning_rate(self):
        pass

    def train(self):
        """Make models eval mode during test time"""
        for name in self.model_names:
            if isinstance(name, str):
                net = getattr(self, 'net' + name)
                net.train()

    def inference(self):
        """Forward function used in test time.

        This function wraps <forward> function in no_grad() so we don't save intermediate steps for backprop
        It also calls <compute_visuals> to produce additional visualization results
        """
        with torch.no_grad():
            self.forward()

    def get_current_losses(self):
        """Return traning losses / errors. train.py will print out these errors on console, and save them to a file"""
        if len(self.loss_names) == 0:
            raise RuntimeError("The list of losses is empty!")
        errors_ret = OrderedDict()
        for name in self.loss_names:
            if isinstance(name, str):
                errors_ret[name] = float(
                    getattr(self,  name))  # float(...) works for both scalar tensor and float number
        return errors_ret

